#include "xsCustomSkinCmd.h"

#include <maya/MSyntax.h>
#include <maya/MArgParser.h>
#include <maya/MGlobal.h>
#include <maya/MArgList.h>
#include <maya/MArgDatabase.h>
#include <maya/MDataHandle.h>
#include <maya/MSelectionList.h>
#include <maya/MFnDependencyNode.h>
#include <maya/MPlug.h>
#include <maya/MFnDagNode.h>
#include <maya/MBoundingBox.h>
#include <maya/MItMeshVertex.h>
#include <maya/MFnTransform.h>
#include <maya/MMatrix.h>
#include <maya/MFnNumericAttribute.h>


struct FlagArgStruct
{
	MString flagShortName;
	MString flagLongName;
	MSyntax::MArgType argType;

	FlagArgStruct(MString sn, MString ln, MSyntax::MArgType type):
	flagShortName(sn), flagLongName(ln), argType(type)
	{};
};

FlagArgStruct argJoint("-j", "-joint", MSyntax::kString);
FlagArgStruct argMesh("-m", "-mesh", MSyntax::kString);


MString xsCustomSkinCmd::cmdName = "createXSCustomSkin";


xsCustomSkinCmd::xsCustomSkinCmd()
{
}

xsCustomSkinCmd::~xsCustomSkinCmd()
{
}

MStatus xsCustomSkinCmd::doIt(const MArgList& args)
{
	MStatus status;
	dagpathJoints.clear();
	dagpathMesh = MDagPath();

	auto AddJoints = [this](MSelectionList& mSelList)
	{
		for (unsigned int i = 0; i < mSelList.length() - 1; ++i)
		{
			MDagPath dagpathJnt;
			mSelList.getDagPath(i, dagpathJnt);
			if (dagpathJnt.isValid() && dagpathJnt.apiType() == MFn::kJoint)
			{
				dagpathJoints.append(dagpathJnt);
			}
			else
			{
				MGlobal::displayError(MString("Influence objects must be joint"));
				return MStatus::kFailure;
			}
		}

		return MStatus::kSuccess;
	};

	auto AddMesh = [this](MSelectionList& mSelList)
	{
		mSelList.getDagPath(mSelList.length() - 1, dagpathMesh);
		MDagPath temp = dagpathMesh;
		temp.extendToShape();
		if (temp.apiType() != MFn::kMesh)
		{
			MGlobal::displayError(MString("Skined object must be mesh"));
			return MStatus::kFailure;
		}

		return MStatus::kSuccess;
	};

	MArgDatabase argData(syntax(), args);
	
	// parser flags
	const char* flagJoint = argJoint.flagShortName.asChar();
	const char* flagMesh = argMesh.flagShortName.asChar();
	if (argData.isFlagSet(flagJoint) and argData.isFlagSet(flagMesh))
	{
		MSelectionList mSelList;
		// parser flag '-j'
		unsigned int jntnum = argData.numberOfFlagUses(flagJoint);
		for (unsigned int i = 0; i < jntnum; ++i)
		{
			MArgList argList;
			argData.getFlagArgumentList(flagJoint, i, argList);

			MString jntName = argList.asString(0);
			mSelList.add(jntName);
		}

		// parser flag '-m'
		MString meshTransform;
		argData.getFlagArgument(flagMesh, 0, meshTransform);
		mSelList.add(meshTransform);

		if (AddJoints(mSelList) == MStatus::kFailure)
			return MStatus::kFailure;

		if (AddMesh(mSelList) == MStatus::kFailure)
			return MStatus::kFailure;
	}
	else
	{
		// if not has flag '-j' '-m'
		
		MSelectionList selectionList;
		argData.getObjects(selectionList);

		if (AddJoints(selectionList) == MStatus::kFailure)
			return MStatus::kFailure;

		if (AddMesh(selectionList) == MStatus::kFailure)
			return MStatus::kFailure;
	}
	
	return AddDeformerAndConnect();
}

MStatus xsCustomSkinCmd::undoIt()
{
	dgModifier.undoIt();
	
	return MStatus::kSuccess;
}

MStatus xsCustomSkinCmd::redoIt()
{
	return MStatus::kFailure;
}

bool xsCustomSkinCmd::isUndoable() const
{
	return true;
}

void* xsCustomSkinCmd::creator()
{
	return new xsCustomSkinCmd();
}

MSyntax xsCustomSkinCmd::createSyntax()
{
	MSyntax syntax;

	syntax.addFlag(argJoint.flagShortName.asChar(), argJoint.flagLongName.asChar(), argJoint.argType);
	syntax.makeFlagMultiUse(argJoint.flagShortName.asChar());

	syntax.addFlag(argMesh.flagShortName.asChar(), argMesh.flagLongName.asChar(), argMesh.argType);

	syntax.setObjectType(MSyntax::kSelectionList);
	syntax.useSelectionAsDefault(true);
	
	return syntax;
}

MStatus xsCustomSkinCmd::AddDeformerAndConnect()
{
	if ((dagpathJoints.length() <= 0) || !dagpathMesh.isValid())
		return MStatus::kFailure;

	MString meshName = dagpathMesh.fullPathName();
	MString pyCmd = "import maya.cmds as cmds\n";
	// lock transform
	MString lockTx = "cmds.setAttr('" + meshName + ".tx" + "', lock=True)\n";
	MString lockTy = "cmds.setAttr('" + meshName + ".ty" + "', lock=True)\n";
	MString lockTz = "cmds.setAttr('" + meshName + ".tz" + "', lock=True)\n";
	MString lockRx = "cmds.setAttr('" + meshName + ".rx" + "', lock=True)\n";
	MString lockRy = "cmds.setAttr('" + meshName + ".ry" + "', lock=True)\n";
	MString lockRz = "cmds.setAttr('" + meshName + ".rz" + "', lock=True)\n";
	MString lockSx = "cmds.setAttr('" + meshName + ".sx" + "', lock=True)\n";
	MString lockSy = "cmds.setAttr('" + meshName + ".sy" + "', lock=True)\n";
	MString lockSz = "cmds.setAttr('" + meshName + ".sz" + "', lock=True)\n";
	pyCmd = pyCmd + lockTx + lockTy + lockTz + lockRx + lockRy + lockRz + lockSx + lockSy + lockSz;
	MGlobal::executePythonCommand(pyCmd, false, true);
	
	// create deformer
	MString cmdAddDeformer = "cmds.deformer('" + meshName + "', type='xsCustomSkin')\n";
	MGlobal::executePythonCommand(cmdAddDeformer, false, true);
	
	// find skin deformer node
	MDagPath tempMesh = dagpathMesh;
	tempMesh.extendToShape();
	MFnDependencyNode mfndpMesh(tempMesh.node());
	MPlug plugInMesh = mfndpMesh.findPlug("inMesh", false);
	MFnDependencyNode mfndpDeformer(plugInMesh.source().node());
	
	// set envelope off
	MPlug plugEnvelope = mfndpDeformer.findPlug("envelope", false);
	plugEnvelope.setFloat(0.0);

	// store geometry world matrix
	MPlug plugPreBindGeomMatrix = mfndpDeformer.findPlug("preBindGeomMatrix", false);
	MDataHandle preBindGeomMatrixHandle = plugPreBindGeomMatrix.asMDataHandle();
	preBindGeomMatrixHandle.setMMatrix(dagpathMesh.inclusiveMatrix());
	plugPreBindGeomMatrix.setValue(preBindGeomMatrixHandle);
	
	// store pre bind matrix and connect attr
	MFnDependencyNode mfndpJnt;
	MPlug plugInfluenceMatrices = mfndpDeformer.findPlug("influenceMatrices", false);
	for (unsigned int i = 0; i < dagpathJoints.length(); ++i)
	{
		MObject obj = dagpathJoints[i].node();
		mfndpJnt.setObject(obj);

		// store pre bind matrix
		MPlug plugWorldMatrixInverse = mfndpJnt.findPlug("worldInverseMatrix", false).elementByLogicalIndex(0);
		StorePreBindMatrix(mfndpDeformer, plugWorldMatrixInverse, i);

		// connect attr
		MPlug plugWorldMatrix = mfndpJnt.findPlug("worldMatrix", false).elementByLogicalIndex(0);
		dgModifier.connect(plugWorldMatrix, plugInfluenceMatrices.elementByLogicalIndex(i));

		// add attr liw
		MFnNumericAttribute mfnNumAttr;
		MObject attrLiw = mfnNumAttr.create("lockInfluenceWeights", "liw", MFnNumericData::kBoolean);
		mfnNumAttr.setKeyable(false);
		mfnNumAttr.setHidden(true);
		mfnNumAttr.setDefault(0);

		mfndpJnt.addAttribute(attrLiw);
	}
	dgModifier.doIt();

	// calculate weights
	CalculateWeights(mfndpDeformer);

	// set envelope on
	plugEnvelope.setFloat(1.0);
	
	return MStatus::kSuccess;
}

void xsCustomSkinCmd::StorePreBindMatrix(MFnDependencyNode& deformerNode, MPlug& worldMatrix, unsigned int index)
{
	MPlug preBindMatrix = deformerNode.findPlug("preBindMatrices", false).elementByLogicalIndex(index);
	
	preBindMatrix.setValue(worldMatrix.asMDataHandle());
}

void xsCustomSkinCmd::CalculateWeights(MFnDependencyNode& deformerNode)
{
	MFnDagNode mfnDagMesh(dagpathMesh);
	MBoundingBox boundbox = mfnDagMesh.boundingBox();
	double diagonal_length_half = (boundbox.max() - boundbox.min()).length() * 0.5;

	MFnTransform mfnTransformJnt;
	MPlug weightList = deformerNode.findPlug("weightList", false);

	MDagPath tempMesh = dagpathMesh;
	tempMesh.extendToShape();
	MItMeshVertex mitVtx(tempMesh);
	mitVtx.reset();

	while (!mitVtx.isDone())
	{
		int index = mitVtx.index();
		MPoint ppos = mitVtx.position(MSpace::kWorld);
		MPlug plugWs = weightList.elementByLogicalIndex(index).child(0);
		double totlaWeights = 0.0;
		MDoubleArray weightList;

		// joint range [0, length-1]
		for (unsigned int j = 0; j < dagpathJoints.length(); ++j)
		{
			mfnTransformJnt.setObject(dagpathJoints[j]);
			MPoint jntPos = mfnTransformJnt.getTranslation(MSpace::kWorld);

			double distanceVtxToJnt = (ppos - jntPos).length();

			double weight = 0.0;
			if (distanceVtxToJnt <= diagonal_length_half)
			{
				// gaus function
				// range [0, 6]
				double x = distanceVtxToJnt / diagonal_length_half * 6.0;

				double a = 2.34;
				double e = 2.71828;
				double b = 3.8 * 1/(sqrt(a) * sqrt(2.0 * 3.14159));
				double exp = -(x * x) / (2.0 * a);

				weight = b * pow(e, exp);
			}

			totlaWeights += weight;
			weightList.append(weight);
		}

		double weightsTotal = 0.0;
		if (totlaWeights > 0.0)
		{
			for (unsigned int j = 0; j < dagpathJoints.length(); ++j)
			{
				double weight = weightList[j] / totlaWeights;
				MPlug plugWj = plugWs.elementByLogicalIndex(j);
				plugWj.setDouble(weight);

				weightsTotal += weight;
			}
		}
		else
		{
			for (unsigned int j = 0; j < dagpathJoints.length(); ++j)
			{
				MPlug plugWj = plugWs.elementByLogicalIndex(j);
				plugWj.setDouble(0);
			}
		}
		
		if (weightsTotal < 1.0)
		{
			MPlug plugWj = plugWs.elementByLogicalIndex(dagpathJoints.length() - 1);
			double existingWeight = plugWj.asDouble();
			plugWj.setDouble(1.0 - weightsTotal + existingWeight);
		}
		
		mitVtx.next();
	}
	
}
